import sets
import copy
import sys

class label_store:
    def __init__(self):
        self.id_dict = {}
        self.label_dict = {}
        self.labels = 0
        
    def add(self,label_sent):
        if self.label_dict.has_key(label_sent):
            return self.label_dict[label_sent]
        else:
            self.labels += 1
            self.label_dict[label_sent] = self.labels
            self.id_dict[self.labels] = label_sent
            return self.label_dict[label_sent]
            
    def finish(self):
        self.label_dict = {}
        
    def get(self,id_sent):
        return self.id_dict[id_sent]

class vertex:

    def __init__(self,id_sent,label_sent):
        self.id = id_sent
        self.label = label_sent
        self.edge_set = sets.Set([])

    def add_edge(self,edge_set_sent):
        self.edge_set = self.edge_set.union(edge_set_sent)

    def debug_print(self,labs_store):
        print "\t\tv " + str(self.id) + " " + labs_store.get(self.label)

class edge:

    def __init__(self,id_sent,label_sent,directed_sent,vertex1_sent,vertex2_sent):
        self.id = id_sent
        self.label = label_sent
        if directed_sent == True:
            self.directed = True
            self.vertex_from = vertex1_sent
            self.vertex_to = vertex2_sent
        else:
            self.directed = False
            self.vertices = sets.Set([vertex1_sent,vertex2_sent])

    def get_vertices(self):
        if self.directed == True:
            return sets.Set([self.vertex_from,self.vertex_to])
        else:
            return self.vertices
        
    def debug_print(self,lab_store):
        if self.directed == True:
            print "\t\td " + str(self.id) + " " + str(self.vertex_from) + " " + str(self.vertex_to) + " " + lab_store.get(self.label)
        else:
            vertices = list(self.vertices)
            print "\t\tu " + str(self.id) + " " + str(vertices[0]) + " " + str(vertices[1]) + " " + lab_store.get(self.label)

class graph:
    
    def __init__(self):
        self.db = {}
        self.vertices = 0
        self.edges = 0

    def add_vertex(self, vertex_id, label_sent):
        vertex_sent =vertex(vertex_id,label_sent)
        self.db["v_" + str(vertex_id)] = vertex_sent
        self.vertices += 1

    def add_edge(self, edge_id,label_sent,directed_sent,vertex1_sent,vertex2_sent):
        edge_sent = edge(edge_id,label_sent,directed_sent,vertex1_sent,vertex2_sent)
        self.db["e_" + str(edge_id)] = edge_sent
        if edge_sent.directed == True:
            vertex = self.get_vertex(edge_sent.vertex_from)
            vertex.add_edge(sets.Set([edge_id]))
            self.db["v_" + str(vertex.id)] = vertex
            vertex = self.get_vertex(edge_sent.vertex_to)
            vertex.add_edge(sets.Set([edge_id]))
            self.db["v_" + str(vertex.id)] = vertex
        else:
            vertices = list(edge_sent.vertices)
            vertex = self.get_vertex(vertices[0])
            vertex.add_edge(sets.Set([edge_id]))
            self.db["v_" + str(vertex.id)] = vertex
            vertex = self.get_vertex(vertices[1])
            vertex.add_edge(sets.Set([edge_id]))
            self.db["v_" + str(vertex.id)] = vertex
        self.edges += 1

    def get_vertex(self,vertex_id):
        return self.db["v_" + str(vertex_id)]

    def get_edge(self,edge_id):
        return self.db["e_" + str(edge_id)]

    def get_vertex_ids(self):
        return xrange(1,self.vertices+1)
    
    def get_edge_ids(self):
        return xrange(1,self.edges+1)
    
    def debug_print(self,labs_store):
        for curr_vertex_id in xrange(1,self.vertices+1):
            curr_vertex = self.get_vertex(curr_vertex_id)
            curr_vertex.debug_print(labs_store)
            
        for curr_edge_id in xrange(1,self.edges+1):
            curr_edge = self.get_edge(curr_edge_id)
            curr_edge.debug_print(labs_store)

    def clear(self):
        self.vertices = 0
        self.edges = 0
        for key in self.db:
            del key
        
class instance:

    def __init__(self, id_sent, ancestor_sent, class_label_sent, vertex_set_sent, edge_set_sent):
        self.id = id_sent
        self.ancestor = ancestor_sent
        self.class_label = class_label_sent
        self.vertex_set = vertex_set_sent
        self.edge_set = edge_set_sent

    def add_vertex(self,vertex_set_sent):
        self.vertex_set = self.vertex_set.union(vertex_set_sent)

    def add_edge(self,edge_set_sent):
        self.edge_set = self.edge_set.union(edge_set_sent)

    def debug_print(self):
        print str(self.id) + " ",
        print str(self.ancestor) + " ",
        print self.class_label + "  v",

        for vertex in self.vertex_set:
            print str(vertex) + " ",
        print " e ",
        for edge in self.edge_set:
            print str(edge) + " ",
        print ""

class instance_set:

    def __init__(self):
        self.instance_dict = {}
        self.instance_id_set = sets.Set([])

    def add(self,instance_list_sent):
        for instance in instance_list_sent:
            original_len = len(self.instance_id_set)
            self.instance_id_set = self.instance_id_set.union([instance.id])
            new_len = len(self.instance_id_set)
            if original_len < new_len:
                self.instance_dict[instance.id] = instance
                
    def remove(self,instance_id_set_sent):
        to_remove = self.instance_id_set.intersection(instance_id_set_sent)
        self.instance_id_set = self.instance_id_set.difference(instance_set_sent.instance_id_set)
        for instance_id in to_remove:
            del self.instance_dict[instance_id]
            
    def get(self,instance_id_sent):
        if self.instance_dict.has_key(instance_id_sent):
            return self.instance_dict[instance_id_sent]
 
    def debug_print(self):
        for instance_id in self.instance_id_set:
            (self.instance_dict[instance_id]).debug_print()

class site_info:

    def __init__(self):
        self.positive = sets.Set([])
        self.negative = sets.Set([])
        self.unknown = sets.Set([])

    def add(self,site_id,class_label_sent):
        if class_label_sent == "xp":
            self.positive = self.positive.union([site_id])
        elif class_label_sent == "xn":
            self.negative = self.negative.union([site_id])
        elif class_label_sent == "xu":
            self.unknown = self.unknown.union([site_id])

    def evaluate(self,covered_set_sent):
        positive_covered = len(self.positive.intersection(covered_set_sent))
        negative_covered = len(self.negative.intersection(covered_set_sent))
        return (float(positive_covered) + float(len(self.negative)) - float(negative_covered))/(float(len(self.positive)) + float(len(self.negative)))

    def covered(self,covered_set_sent):
        positive_covered = self.positive.intersection(covered_set_sent)
        negative_covered = self.negative.intersection(covered_set_sent)
        unknown_covered = self.unknown.intersection(covered_set_sent)
        return (positive_covered,negative_covered,unknown_covered)

    def remove(self,covered_set_sent):
        self.positive = self.positive.difference(covered_set_sent)
        self.unknown = self.unknown.difference(covered_set_sent)

    def debug_print(self):
        print "Positive Sites: ",
        for site in self.positive:
            print site,
            print " ",

        print ""

        print "Positive Sites: ",
        for site in self.negative:
            print site,
            print " ",

        print ""

        print "Positive Sites: ",
        for site in self.unknown:
            print site,
            print " "

        print ""
   
def create_normalized_graph(instance_sent,base_graph):
    normalized_graph = graph()
    normalized_vertices = {}
    normalized_edges = {}
    vertex_list = list(instance_sent.vertex_set)
    edge_list = list(instance_sent.edge_set)
    for normalized_vertex_id in xrange(1,len(instance_sent.vertex_set)+1):
        normalized_vertices[vertex_list[normalized_vertex_id-1]] = normalized_vertex_id
    for normalized_edge_id in xrange(1,len(instance_sent.edge_set)+1):
        normalized_edges[edge_list[normalized_edge_id-1]] = normalized_edge_id
    for vertex_id in instance_sent.vertex_set:
        base_vertex = base_graph.get_vertex(vertex_id)
        normalized_graph.add_vertex(normalized_vertices[vertex_id],base_vertex.label)
    for edge_id in instance_sent.edge_set:
        base_edge = base_graph.get_edge(edge_id)
        base_edge_vertices = list(base_edge.get_vertices())
        normalized_graph.add_edge(normalized_edges[edge_id],base_edge.label,base_edge.directed,normalized_vertices[base_edge_vertices[0]],\
                               normalized_vertices[base_edge_vertices[1]])
    return normalized_graph        


def build_vertex_dict(graph_sent):
    vertex_dict = {}

    for curr_vertex_id in xrange(1,graph_sent.vertices+1):
        curr_vertex = graph_sent.get_vertex(curr_vertex_id)
        curr_vertex_edge_info = []
        curr_vertex_edge_info.append((-1,-1,-1,curr_vertex.label))
        for curr_edge_id in curr_vertex.edge_set:
            curr_edge = graph_sent.get_edge(curr_edge_id)
            if curr_edge.directed == True:
                if curr_edge.vertex_from == curr_vertex_id:
                    curr_vertex_edge_info.append((curr_edge.label,(graph_sent.get_vertex(curr_edge.vertex_to)).label,1,1))
                else:
                    curr_vertex_edge_info.append((curr_edge.label,(graph_sent.get_vertex(curr_edge.vertex_from)).label,1,0))
            else:
                curr_vertex_edge_info.append((curr_edge.label, (graph_sent.get_vertex(curr_edge.vertices.difference([curr_vertex_id]).pop())).label,0,0))
        curr_vertex_edge_info.sort()
        curr_signature = str(curr_vertex_edge_info)
        if vertex_dict.has_key(curr_signature):
            vertex_dict[curr_signature] = vertex_dict[curr_signature].union([curr_vertex_id])
        else:
            vertex_dict[curr_signature] = sets.Set([curr_vertex_id])
    return vertex_dict

def read_graph(graph_file_name):
    graph_sent = graph()
    curr_labs = label_store()
    try:
        graph_file = open(graph_file_name,"r")
    except:
        return -1
    for line in graph_file:
        words = line.split()
        if words[0] == "v":
            graph_sent.add_vertex(int(words[1]), curr_labs.add(words[2]))
        elif words[0] == "u":
            graph_sent.add_edge(int(words[1]), curr_labs.add(words[4]), False, int(words[2]), int(words[3]))
        elif words[0] == "d":
            graph_sent.add_edge(int(words[1]), curr_labs.add(words[4]), True, int(words[2]), int(words[3]))
    graph_file.close()
    curr_labs.finish()
    return (graph_sent,curr_labs)

def read_instances(instance_file_name):
    try:
        instance_file = open(instance_file_name,"r")
    except:
        return -1
    curr_instance_set = instance_set()
    curr_site_info = site_info()
    for line in instance_file:
        words = line.split()
        curr_instance_id = int(words[0])
        curr_instance_class_label = words[1]
        curr_site_info.add(curr_instance_id,curr_instance_class_label)
        curr_vertex_set = sets.Set([])
        curr_edge_set = sets.Set([])
        for word_index in xrange(3,len(words)):
            if words[word_index] == "e":
                edge_begin = word_index + 1
                break
            else:
                curr_vertex_set = curr_vertex_set.union([int(words[word_index])])
        for word_index in xrange(edge_begin,len(words)):
            curr_edge_set = curr_edge_set.union([int(words[word_index])])
        new_instance = instance(curr_instance_id,curr_instance_id,curr_instance_class_label,curr_vertex_set,curr_edge_set)    
        curr_instance_set.add([new_instance])
    
    return (curr_site_info,curr_instance_set)
        
def expand_instance_set(instance_set_sent,base_graph):
    expanded_instance_set = instance_set()
    instance_count = 0
    for curr_instance_id in instance_set_sent.instance_id_set:
        curr_instance = instance_set_sent.get(curr_instance_id)
        outgoing_edges = sets.Set([])
        for curr_vertex in curr_instance.vertex_set:
            outgoing_edges = outgoing_edges.union((base_graph.get_vertex(curr_vertex)).edge_set)
        expansion_edges = outgoing_edges.difference(curr_instance.edge_set)
     
        for new_edge in expansion_edges:
            new_edges = curr_instance.edge_set.union([new_edge])
            new_edge_vertices = (base_graph.get_edge(new_edge)).get_vertices()
            new_vertices = new_edge_vertices.union(curr_instance.vertex_set)
            instance_count += 1
            new_instance = instance(instance_count,curr_instance.ancestor,curr_instance.class_label,new_vertices,new_edges)
            expanded_instance_set.add([new_instance])
    return expanded_instance_set

def group_instances(instance_set_sent,base_graph,site_info_sent,lab_store):
    instance_group_list = []
    original_instance_id_set = instance_set_sent.instance_id_set
    while len(original_instance_id_set) >= 2:
        curr_instance_id = original_instance_id_set.pop()
        remaining_instance_id_set = original_instance_id_set
        curr_instance_group = sets.Set([curr_instance_id])
        curr_instance = instance_set_sent.get(curr_instance_id)
        curr_ancestor = curr_instance.ancestor
        curr_instance_ancestor_group = sets.Set([curr_ancestor])

        
        curr_graph = create_normalized_graph(curr_instance,base_graph)
        curr_vertex_dict = build_vertex_dict(curr_graph)
        
        for to_compare_instance_id in remaining_instance_id_set:
            to_compare_instance = instance_set_sent.get(to_compare_instance_id)
            to_compare_ancestor = to_compare_instance.ancestor
            to_compare_graph = create_normalized_graph(to_compare_instance,base_graph)
            to_compare_vertex_dict = build_vertex_dict(to_compare_graph)
            if graph_isomorphism(curr_graph,to_compare_graph,curr_vertex_dict,to_compare_vertex_dict) == True:
                curr_instance_group = curr_instance_group.union([to_compare_instance_id])
                curr_instance_ancestor_group = curr_instance_ancestor_group.union([to_compare_ancestor])

        
        curr_value = site_info_sent.evaluate(curr_instance_ancestor_group)
        
        original_instance_id_set = original_instance_id_set.difference(curr_instance_group)
        
        instance_group_list.append((curr_value, curr_graph, curr_instance_ancestor_group, curr_instance_group))

    for info in instance_group_list:
        print "Considering..."
        print "\tPattern:"
        (info[1]).debug_print(lab_store)
        print "\tValue: " + str(info[0])
    return instance_group_list

def prune(instance_group_list, beam):
    instance_group_list.sort()
    instance_group_list_size = len(instance_group_list)
    if instance_group_list_size < beam:
        retain = instance_group_list_size
    else:
        retain = beam
    pruned_instance_group_list = []
    for retained_instance_group in xrange(1,retain+1):
        pruned_instance_group_list.append(instance_group_list.pop())
    return pruned_instance_group_list

def ungroup_instances(pruned_instance_group_list,instance_set_sent):
    instance_group_list_sent = []
    for group in pruned_instance_group_list:
        instance_group_list_sent.append(group[3])
        
    ungrouped_instance_id_set = sets.Set([])
    for instance_group in instance_group_list_sent:
        ungrouped_instance_id_set = ungrouped_instance_id_set.union(instance_group)

    ungrouped_instance_set = instance_set()
    for curr_instance_id in ungrouped_instance_id_set:
        curr_instance = instance_set_sent.get(curr_instance_id)
        ungrouped_instance_set.add([curr_instance])

    return ungrouped_instance_set

class iteration:

    def __init__(self,id_sent):
        self.id = id_sent
        self.best_value = 0
        self.instance_ancestor_group = []

    def add_result(self,value_sent, graph_sent, instance_ancestor_group_sent):
        if self.best_value < value_sent:
            self.best_value = value_sent
            self.graph = graph_sent
            self.instance_ancestor_group = instance_ancestor_group_sent
            

    def print_debug(self):
        print "-----------ITERATION: " + str(self.id) + " --------------"
        print "VALUE: " + str(self.best_value)
        print "SITES COVERED: " + str(self.instance_ancestor_group)
        print "PATTERN:"
        self.graph.debug_print()
        
class iteration_info:

    def __init__(self):
        self.iteration_list = []

    def add(self,iteration_sent):
        self.iteration_list.append(iteration_sent)
    
    def debug_print(self):
        for curr_iteration in self.iteration_list:
            curr_iteration.print_debug()

def get_initial_instances(site_info, original_sites):
   curr_sites = instance_set()
   for curr_site_id in site_info.positive:
       curr_site = original_sites.get(curr_site_id)
       curr_sites.add([curr_site])
   for curr_site_id in site_info.negative:
       curr_site = original_sites.get(curr_site_id)
       curr_sites.add([curr_site])
   return curr_sites

class mappings:
    def __init__(self):
        self.map = []

    def clear(self):
        self.map = []

    def get_map(self):
        return self.map
    
    def permute_counts(self, group_counts, per_list):
        if len(group_counts) == 0:
            self.map.append(per_list)
        else:
            for element in xrange(0,(group_counts[:1]).pop()):
                self.permute_counts(group_counts[1:], per_list + [element])
     
    def permute(self, input_set, per_list):
        if len(input_set) == 0:
            self.map.append(per_list)
        else:
            for element in input_set:
                self.permute(input_set.difference([element]), per_list + [element])
            

def generate_mappings(group_list):
    group_permute_list = []
    group_permute_length_list = []
    
    for group in group_list:
        dummy = mappings()
        dummy.permute(group,[])
        group_map = dummy.get_map()
        group_permute_list.append(group_map)
        group_permute_length_list.append(len(group_map))
        dummy.clear()

    dummy.permute_counts(group_permute_length_list,[])
    group_mapping = dummy.get_map()
    dummy.clear()

    final_mappings = []
    for curr_mapping in group_mapping:
        curr_map = []
        for group_count in xrange(0,len(group_list)):
            curr_map = curr_map + group_permute_list[group_count][curr_mapping[group_count]]
        final_mappings.append(curr_map)
    return final_mappings

def check_isomorphism_mapping(graph1,graph2,vertex_map):
    vertex_map_dict = {}

    for vertex_id in xrange(0,len(vertex_map[0])):
        vertex_map_dict[vertex_map[0][vertex_id]] = vertex_map[1][vertex_id]

    graph1_edge_list = []
    for curr_edge_id in graph1.get_edge_ids():
        curr_edge = graph1.get_edge(curr_edge_id)
        if curr_edge.directed == True:
             graph1_edge_list.append((curr_edge.label,1,vertex_map_dict[curr_edge.vertex_from], vertex_map_dict[curr_edge.vertex_to]))
        else:
            curr_edge_vertices = list(curr_edge.vertices)
            curr_edge_vertices.sort()
            graph1_edge_list.append((curr_edge.label,0,vertex_map_dict[curr_edge_vertices[0]],vertex_map_dict[curr_edge_vertices[1]]))

    graph1_edge_list.sort()
    graph1_signature = str(graph1_edge_list)
 
    graph2_edge_list = []
    for curr_edge_id in graph2.get_edge_ids():
        curr_edge = graph2.get_edge(curr_edge_id)
        if curr_edge.directed == True:
             graph2_edge_list.append((curr_edge.label,1,curr_edge.vertex_from,curr_edge.vertex_to))
        else:
            curr_edge_vertices = list(curr_edge.vertices)
            curr_edge_vertices.sort()
            graph2_edge_list.append((curr_edge.label,0,curr_edge_vertices[0],curr_edge_vertices[1]))

    graph2_edge_list.sort()
    graph2_signature = str(graph2_edge_list)

    if graph1_signature == graph2_signature:
        return True
    else:
        return False

def graph_isomorphism(graph1,graph2,vertex_dict1 = None,vertex_dict2 = None):

    if vertex_dict1 == None:
        vertex_dict1 = build_vertex_dict(graph1)
    if vertex_dict2 == None:
        vertex_dict2 = build_vertex_dict(graph2)

    vertex_group_list1 = []
    vertex_group_list2 = []

    if len(vertex_dict1.keys()) == len(vertex_dict2.keys()):
        for vertex_signature in vertex_dict1:
            if vertex_dict2.has_key(vertex_signature):
                if len(vertex_dict1[vertex_signature]) == len(vertex_dict2[vertex_signature]):
                    vertex_group_list1 = vertex_group_list1 + list(vertex_dict1[vertex_signature])
                    vertex_group_list2.append(vertex_dict2[vertex_signature])
                else:
                    return False
            else:
                return False
        vertex_mapping_list = generate_mappings(vertex_group_list2)
        return_flag = False
        for vertex_mapping in vertex_mapping_list:
            if check_isomorphism_mapping(graph1,graph2,[vertex_group_list1,vertex_mapping]):
                return True
        return False
    else:
        return False


def print_predictions(iteration_info,site_info):
    print ""
    print "************************* Classifications ***********************"
    print "Pattern\t\t",
    for iteration_id in xrange(1,len(iteration_info.iteration_list) + 1):
        print iteration_id,
        print "\t",
    print ""
    print "-------"
    print "Site"
    site_count = len(site_info.positive) + len(site_info.negative) + len(site_info.unknown)
    for site_id in xrange(1,site_count+1):
        print site_id,
        print "\t\t",
        for iteration in iteration_info.iteration_list:
            if site_id in iteration.instance_ancestor_group.intersection(site_info.positive):
                print "XP\t",
            elif site_id in iteration.instance_ancestor_group.intersection(site_info.negative):
                print "XP\t",
            elif site_id in iteration.instance_ancestor_group.intersection(site_info.unknown):
                print "XP\t",
            else:
                print "XN\t",
        print ""


def print_result(iteration_info,site_info,labs_store):
    print "************************* Current Result ***********************"

    for iteration in iteration_info.iteration_list:
        print "Iteration: " + str(iteration.id) 
        print "\tPattern:"
        iteration.graph.debug_print(labs_store)
        print "\tValue: " + str(iteration.best_value)
        print "\tPositive Covered:",
        print len(iteration.instance_ancestor_group.intersection(site_info.positive))
        print "\tNegative Covered:",
        print len(iteration.instance_ancestor_group.intersection(site_info.negative))
        print "\tUnknown Covered:",
        print len(iteration.instance_ancestor_group.intersection(site_info.unknown))
    
def print_set(set_sent):
    if len(set_sent) > 0:
        for element in set_sent:
            print str(element) + " ",
        print ""
    else:
        print "None"

def subdueMPS(graph_file_name,site_file_name, beam, depth, iterations):
   
    (my_graph, my_labels) = read_graph(graph_file_name)
    (my_site_info, original_sites) = read_instances(site_file_name)
    original_site_info = copy.copy(my_site_info)
    my_iteration_info = iteration_info()
    iteration_count = 1

    print "************************* SubdueMPS ***********************"
    print "Vertices :" + str(my_graph.vertices)
    print "Edges :" + str(my_graph.edges)
    print "Positive Sites :",
    print len(original_site_info.positive)
    print "Negative Sites :",
    print len(original_site_info.negative)
    print "Unknown Sites :",
    print len(original_site_info.unknown)

    
    while iteration_count <= iterations and len(my_site_info.positive) > 0:
        curr_iteration = iteration(iteration_count)
        print ""
        print "*********************** Iteration: " + str(iteration_count) + " *********************"

        my_instance_set = get_initial_instances(my_site_info,original_sites)
        depth_count = depth
        while depth_count > 0:
            my_expanded_instance_set = expand_instance_set(my_instance_set,my_graph)
            
            my_instance_groups = group_instances(my_expanded_instance_set,my_graph,my_site_info,my_labels)
            my_pruned_groups = prune(my_instance_groups,beam)
            (curr_value, curr_graph, curr_instance_ancestor_group,curr_instance_group) = my_pruned_groups[0]
            curr_iteration.add_result(curr_value, curr_graph, curr_instance_ancestor_group)
            my_instance_set = ungroup_instances(my_pruned_groups,my_expanded_instance_set)
            depth_count -= 1

        my_site_info.remove(curr_iteration.instance_ancestor_group)
        
        my_iteration_info.add(curr_iteration)
        iteration_count += 1
        print_result(my_iteration_info,original_site_info,my_labels)

    print_result(my_iteration_info,original_site_info,my_labels)
    print_predictions(my_iteration_info,original_site_info)
    print "****************************************************************"


if __name__ == "__main__":
    subdueMPS("/Users/ketkar/Documents/workspace/subdueMPS/genes.mpsg","/Users/ketkar/Documents/workspace/subdueMPS/cell_growth.s",4,5,5)
